import numpy as np
import pandas as pd
from scipy.optimize import least_squares, minimize
import cv2
import os
from io import StringIO
import re
import matplotlib.pyplot as plt
import tkinter as tk
from tkinter import ttk

# 文件路径设置
CSV_FILE_PATH = '靶标数据_20250527_201524.csv'  # CSV文件路径

# 误差参数设置
PIXEL_ERROR_STD = 0  # 像素坐标标准差（像素）
TARGET_COORD_ERROR_STD = 0  # 靶标坐标系坐标标准差（mm）
TOTAL_STATION_ERROR_STD = 0  # 全站仪测量靶标原点标准差（mm）

# 优化参数设置
RANSAC_ITERATIONS = 100  # RANSAC迭代次数
RANSAC_THRESHOLD = 2.0  # RANSAC阈值（像素）
MIN_INLIER_RATIO = 0.5  # 最小内点比例

def euler_to_matrix(euler):
    """欧拉角转旋转矩阵，ZYX顺序，单位为度"""
    rx, ry, rz = np.radians(euler)
    Rx = np.array([
        [1, 0, 0],
        [0, np.cos(rx), -np.sin(rx)],
        [0, np.sin(rx), np.cos(rx)]
    ])
    Ry = np.array([
        [np.cos(ry), 0, np.sin(ry)],
        [0, 1, 0],
        [-np.sin(ry), 0, np.cos(ry)]
    ])
    Rz = np.array([
        [np.cos(rz), -np.sin(rz), 0],
        [np.sin(rz), np.cos(rz), 0],
        [0, 0, 1]
    ])
    return Rz @ Ry @ Rx

def read_simulation_data(file_path=CSV_FILE_PATH):
    """读取仿真数据CSV文件并打印信息"""
    # 获取当前脚本所在目录
    current_dir = os.path.dirname(os.path.abspath(__file__))
    full_path = os.path.join(current_dir, file_path)
    if not os.path.exists(full_path):
        print(f"错误：找不到文件 '{file_path}'")
        print(f"请确保文件位于以下目录：{current_dir}")
        return None, None, None

    # 读取文件内容
    with open(full_path, 'r', encoding='utf-8') as f:
        lines = f.readlines()

    # 初始化参数字典
    camera_params = {}
    camera_extrinsic = {}
    df = None
    current_section = None
    data_lines = []

    # 逐行处理文件
    for line in lines:
        line = line.strip()
        if not line:
            continue

        # 检查是否是新的部分
        if line.startswith('==='):
            if '相机参数' in line:
                current_section = 'camera_params'
            elif '相机外参' in line:
                current_section = 'camera_extrinsic'
            elif '靶标数据' in line:
                current_section = 'target_data'
            continue

        # 根据当前部分处理数据
        if current_section == 'camera_params':
            if ',' in line:
                parts = line.split(',')
                if len(parts) == 2:
                    key = parts[0].strip()
                    try:
                        value = float(parts[1].strip())
                        camera_params[key] = value
                    except ValueError:
                        continue

        elif current_section == 'camera_extrinsic':
            if ',' in line and ('相机位置' in line or '相机欧拉角' in line):
                parts = line.split(',')
                if len(parts) == 2:
                    key = parts[0].strip()
                    try:
                        value = float(parts[1].strip())
                        camera_extrinsic[key] = value
                    except ValueError:
                        continue

        elif current_section == 'target_data':
            data_lines.append(line)

    # 处理靶标数据
    if data_lines:
        df = pd.read_csv(StringIO('\n'.join(data_lines)))

    # 打印读取到的参数
    print("\n=== 读取到的相机参数 ===")
    for key, value in sorted(camera_params.items()):
        print(f"{key}: {value}")

    print("\n=== 读取到的相机外参 ===")
    for key, value in sorted(camera_extrinsic.items()):
        print(f"{key}: {value}")

    # 验证必要的参数
    required_params = ['X方向焦距(像素)', 'Y方向焦距(像素)', '主点X坐标(像素)', '主点Y坐标(像素)']
    required_extrinsic = ['相机位置X(mm)', '相机位置Y(mm)', '相机位置Z(mm)',
                        '相机欧拉角X(度)', '相机欧拉角Y(度)', '相机欧拉角Z(度)']

    missing_params = [param for param in required_params if param not in camera_params]
    if missing_params:
        print("\n错误：缺少以下相机参数：")
        for param in missing_params:
            print(f"  - {param}")
        return None, None, None

    missing_extrinsic = [param for param in required_extrinsic if param not in camera_extrinsic]
    if missing_extrinsic:
        print("\n错误：缺少以下相机外参：")
        for param in missing_extrinsic:
            print(f"  - {param}")
        return None, None, None

    if df is None or df.empty:
        print("错误：未找到有效的靶标数据")
        return None, None, None

    return df, camera_params, camera_extrinsic

def matrix_to_euler(R):
    """旋转矩阵转欧拉角（ZYX，度）"""
    sy = np.sqrt(R[0,0]**2 + R[1,0]**2)
    singular = sy < 1e-6
    if not singular:
        x = np.arctan2(R[2,1], R[2,2])
        y = np.arctan2(-R[2,0], sy)
        z = np.arctan2(R[1,0], R[0,0])
    else:
        x = np.arctan2(-R[1,2], R[1,1])
        y = np.arctan2(-R[2,0], sy)
        z = 0
    return np.degrees([x, y, z])

def compute_reprojection_error(params, Otw_list, image_points_list, object_points):
    """计算重投影误差"""
    fx, fy, cx, cy = params[:4]
    euler = params[4:7]
    T_wc = params[7:10]
    R_wc = euler_to_matrix(euler)
    K = np.array([[fx, 0, cx],
                  [0, fy, cy],
                  [0,  0,  1]])
    
    total_error = 0
    for Otw, img_pts in zip(Otw_list, image_points_list):
        # 计算靶标点在相机坐标系下的位置
        P_c = R_wc @ object_points.T + T_wc[:, np.newaxis]
        # 投影到图像平面
        P_2d = K @ P_c
        P_2d = P_2d[:2] / P_2d[2]
        # 计算重投影误差
        error = np.linalg.norm(P_2d.T - img_pts, axis=1)
        total_error += np.mean(error)
    
    return total_error / len(Otw_list)

def ransac_optimization(x0, Otw_list, image_points_list, object_points):
    """使用RANSAC思想进行粗优化"""
    best_params = x0.copy()
    best_error = float('inf')
    best_inliers = None
    
    for _ in range(RANSAC_ITERATIONS):
        # 随机选择一部分点
        n_points = len(Otw_list)
        sample_size = max(3, int(n_points * 0.7))  # 至少使用3个点
        indices = np.random.choice(n_points, sample_size, replace=False)
        
        sample_Otw = [Otw_list[i] for i in indices]
        sample_img_pts = [image_points_list[i] for i in indices]
        
        # 使用选中的点进行优化
        result = least_squares(
            lambda x: residual_func(x, sample_Otw, sample_img_pts, object_points),
            x0,
            method='trf',
            ftol=1e-3,
            xtol=1e-3,
            max_nfev=30
        )
        
        # 计算所有点的重投影误差
        error = compute_reprojection_error(result.x, Otw_list, image_points_list, object_points)
        
        # 找出内点
        inliers = []
        for i, (Otw, img_pts) in enumerate(zip(Otw_list, image_points_list)):
            if compute_reprojection_error(result.x, [Otw], [img_pts], object_points) < RANSAC_THRESHOLD:
                inliers.append(i)
        
        # 更新最佳结果
        if len(inliers) / n_points > MIN_INLIER_RATIO and error < best_error:
            best_params = result.x.copy()
            best_error = error
            best_inliers = inliers
    
    return best_params, best_inliers

def residual_func(params, Otw_list, image_points_list, object_points, iter_log=None):
    """计算残差函数"""
    fx, fy, cx, cy = params[:4]
    euler = params[4:7]
    T_wc = params[7:10]
    R_wc = euler_to_matrix(euler)
    K = np.array([[fx, 0, cx],
                  [0, fy, cy],
                  [0,  0,  1]])
    
    res = []
    for Otw, img_pts in zip(Otw_list, image_points_list):
        # PnP解靶标原点在相机坐标系下的位置
        _, rvec, tvec = cv2.solvePnP(object_points, img_pts, K, None, flags=cv2.SOLVEPNP_ITERATIVE)
        T_tc = tvec.flatten()
        # 公式2-1: Otc = R_wc * Otw + T_wc
        Otc = R_wc @ Otw + T_wc
        # 残差 = Otc - T_tc
        res.extend(Otc - T_tc)
    
    res = np.array(res)
    
    # 直接返回残差
    return res

def multi_stage_optimization(x0, Otw_list, image_points_list, object_points):
    """多阶段优化策略"""
    iter_log = []
    
    result = least_squares(
        lambda x: residual_func(x, Otw_list, image_points_list, object_points, iter_log),
        x0,
        method='trf',
        ftol=1e-8,
        xtol=1e-8,
        max_nfev=100,
        verbose=2
    )
    
    return result, iter_log

def parse_pixel_coords(pixel_str):
    """解析像素坐标字符串，并添加随机误差"""
    coords = re.findall(r'\(([\d\.\-]+),\s*([\d\.\-]+)\)', pixel_str)
    coords = np.array([[float(x), float(y)] for x, y in coords], dtype=np.float32)
    # 添加像素坐标误差
    coords += np.random.normal(0, PIXEL_ERROR_STD, coords.shape)
    return coords

def parse_vector(vec_str):
    """解析向量字符串，并添加随机误差"""
    nums = vec_str.strip('[]').split(',')
    vec = np.array([float(x) for x in nums])
    # 添加全站仪测量误差
    vec += np.random.normal(0, TOTAL_STATION_ERROR_STD, vec.shape)
    return vec

def main():
    # 读取数据获取真值
    df, camera_params, camera_extrinsic = read_simulation_data(CSV_FILE_PATH)

    if df is None or camera_params is None or camera_extrinsic is None:
        print("错误：无法从CSV文件读取参数")
        exit(1)

    # 从CSV文件中读取的真值
    fx_true = camera_params['X方向焦距(像素)']
    fy_true = camera_params['Y方向焦距(像素)']
    cx_true = camera_params['主点X坐标(像素)']
    cy_true = camera_params['主点Y坐标(像素)']

    # 从CSV文件中读取相机外参真值
    T_cw_true = np.array([
        camera_extrinsic['相机位置X(mm)'],
        camera_extrinsic['相机位置Y(mm)'],
        camera_extrinsic['相机位置Z(mm)']
    ])

    # 从欧拉角计算旋转矩阵
    euler_true = np.array([
        camera_extrinsic['相机欧拉角X(度)'],
        camera_extrinsic['相机欧拉角Y(度)'],
        camera_extrinsic['相机欧拉角Z(度)']
    ])

    R_cw_true = euler_to_matrix(euler_true)

    R_wc_true = R_cw_true.T
    T_wc_true = -R_wc_true @ T_cw_true

    # 添加扰动
    np.random.seed(42)
    fx0 = fx_true + np.random.normal(0, 100)   # 焦距fx初值，均值为真实值，标准差100
    fy0 = fy_true + np.random.normal(0, 100)   # 焦距fy初值，均值为真实值，标准差100
    cx0 = cx_true + np.random.normal(0, 10)    # 主点cx初值，均值为真实值，标准差10
    cy0 = cy_true + np.random.normal(0, 10)    # 主点cy初值，均值为真实值，标准差10
    # 欧拉角扰动
    euler0 = euler_true + np.random.normal(0, 10, 3)  # 3个角度扰动（度），标准差10度
    R_cw0 = cv2.Rodrigues(np.radians(euler0))[0]
    T_cw0 = T_cw_true + np.random.normal(0, 50, 3)
    R_wc0 = R_cw0.T
    T_wc0 = -R_wc0 @ T_cw0
    # 组装优化所需数据
    Otw_list = []
    image_points_list = []
    for _, row in df.iterrows():
        img_pts = parse_pixel_coords(row['靶标点像素坐标'])
        image_points_list.append(img_pts)
        Otw = parse_vector(row['靶标原点在全站仪坐标系下的坐标Ot_c'])
        Otw_list.append(Otw)
    # 靶标三维点（单位mm），需与你仿真时一致
    object_points = np.array([
        [0,   0,   0],    # P1
        [500, 0,   0],    # P2
        [500, 500, 0],    # P3
        [0,   500, 0],    # P4
        [-250,   250,   -300],  # P5
        [750, 250, -300],  # P6
    ], dtype=np.float32)

    # 添加靶标坐标系坐标误差
    object_points += np.random.normal(0, TARGET_COORD_ERROR_STD, object_points.shape)

    # 初值
    x0 = np.array([fx0, fy0, cx0, cy0, *np.degrees(cv2.Rodrigues(R_wc0)[0].flatten()), *T_wc0])

    # 使用多阶段优化
    result, iter_log = multi_stage_optimization(x0, Otw_list, image_points_list, object_points)

    fx_opt = result.x[0]
    fy_opt = result.x[1]
    cx_opt = result.x[2]
    cy_opt = result.x[3]
    T_wc_opt = result.x[7:10]
    euler_opt = result.x[4:7]
    R_wc_opt = euler_to_matrix(euler_opt)
    R_cw_opt = R_wc_opt.T
    T_cw_opt = -R_cw_opt @ T_wc_opt
    euler_opt_cw = matrix_to_euler(R_cw_opt)


    # 输出对比信息
    print("\n=== 优化结果与真值对比 ===")
    print("{:<16s} {:>16s} {:>16s} {:>16s}".format("参数", "优化值", "真值", "差值"))
    print("-" * 64)
    print("{:<16s} {:>16.8f} {:>16.8f} {:>16.8f}".format("fx", fx_opt, fx_true, fx_opt - fx_true))
    print("{:<16s} {:>16.8f} {:>16.8f} {:>16.8f}".format("fy", fy_opt, fy_true, fy_opt - fy_true))
    print("{:<16s} {:>16.8f} {:>16.8f} {:>16.8f}".format("cx", cx_opt, cx_true, cx_opt - cx_true))
    print("{:<16s} {:>16.8f} {:>16.8f} {:>16.8f}".format("cy", cy_opt, cy_true, cy_opt - cy_true))
    print("{:<16s} {:>16.8f} {:>16.8f} {:>16.8f}".format("T_wc[0]", T_cw_opt[0], T_cw_true[0], T_wc_opt[0] - T_wc_true[0]))
    print("{:<16s} {:>16.8f} {:>16.8f} {:>16.8f}".format("T_wc[1]", T_cw_opt[1], T_cw_true[1], T_wc_opt[1] - T_wc_true[1]))
    print("{:<16s} {:>16.8f} {:>16.8f} {:>16.8f}".format("T_wc[2]", T_cw_opt[2], T_cw_true[2], T_wc_opt[2] - T_wc_true[2]))
    print("{:<16s} {:>16.8f} {:>16.8f} {:>16.8f}".format("Euler_X", euler_opt_cw[0], euler_true[0], euler_opt_cw[0] - euler_true[0]))
    print("{:<16s} {:>16.8f} {:>16.8f} {:>16.8f}".format("Euler_Y", euler_opt_cw[1], euler_true[1], euler_opt_cw[1] - euler_true[1]))
    print("{:<16s} {:>16.8f} {:>16.8f} {:>16.8f}".format("Euler_Z", euler_opt_cw[2], euler_true[2], euler_opt_cw[2] - euler_true[2]))
    print("-" * 64)
    
    print("\n旋转矩阵 R_wc (优化值):")
    R_wc_opt
    for row in R_wc_opt:
        print("  [" + ", ".join(f"{v:.8f}" for v in row) + "]")
    print("旋转矩阵 R_wc_true (真值):")
    for row in R_wc_true:
        print("  [" + ", ".join(f"{v:.8f}" for v in row) + "]")


if __name__ == "__main__":
    main()
